package db

import (
	"database/sql"
	"encoding/json"
	"fmt"
	"math"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"sync"
	"time"

	"goVector/pkg/model"

	"github.com/jmoiron/sqlx"
	_ "github.com/mattn/go-sqlite3"
)

// VectorDB 表示向量数据库接口
type VectorDB interface {
	InitDB() error
	AddDocument(content string, vector []float32) (int64, error)
	UpdateDocument(id int64, content string, vector []float32) error
	DeleteDocument(id int64) error
	GetDocument(id int64) (*model.Document, error)
	SearchSimilar(vector []float32, limit int, minScore float32) ([]model.SearchResult, error)
	ListDocuments(limit, offset int) ([]model.Document, int64, error)
	Close() error
	SaveToDisk() error
}

// SQLiteVectorDB 实现了基于SQLite的向量数据库
type SQLiteVectorDB struct {
	db               *sqlx.DB
	filePath         string
	memoryMode       bool
	lastPersistTime  time.Time
	mu               sync.RWMutex
	modified         bool
	autoSaveInterval time.Duration
	stopAutoSave     chan struct{}
}

// CollectionManager 管理多个向量数据库集合
type CollectionManager struct {
	baseDir     string
	collections map[string]*SQLiteVectorDB
	metadata    model.CollectionsMetadata
	metaFile    string
	mu          sync.RWMutex
}

// NewCollectionManager 创建新的集合管理器
func NewCollectionManager(baseDir string) *CollectionManager {
	// 确保基础目录存在
	os.MkdirAll(baseDir, 0755)

	// 初始化集合管理器
	cm := &CollectionManager{
		baseDir:     baseDir,
		collections: make(map[string]*SQLiteVectorDB),
		metaFile:    filepath.Join(baseDir, "collections_metadata.json"),
	}

	// 尝试加载元数据
	if err := cm.loadMetadata(); err != nil {
		fmt.Printf("加载元数据失败，将使用扫描方式: %v\n", err)
		// 如果没有元数据或加载失败，尝试扫描目录
		cm.scanDirectoryForCollections()
	}

	return cm
}

// 加载元数据文件
func (cm *CollectionManager) loadMetadata() error {
	// 检查元数据文件是否存在
	if _, err := os.Stat(cm.metaFile); os.IsNotExist(err) {
		return fmt.Errorf("元数据文件不存在")
	}

	// 读取元数据文件
	data, err := os.ReadFile(cm.metaFile)
	if err != nil {
		return fmt.Errorf("读取元数据文件失败: %w", err)
	}

	// 解析JSON
	if err := json.Unmarshal(data, &cm.metadata); err != nil {
		return fmt.Errorf("解析元数据失败: %w", err)
	}

	// 根据元数据加载集合
	for _, meta := range cm.metadata.Collections {
		// 检查文件是否存在
		if _, err := os.Stat(meta.FilePath); os.IsNotExist(err) {
			fmt.Printf("警告: 集合文件不存在: %s\n", meta.FilePath)
			continue
		}

		// 创建数据库连接
		db, err := NewSQLiteVectorDB(meta.FilePath)
		if err != nil {
			fmt.Printf("加载集合 %s 失败: %v\n", meta.Name, err)
			continue
		}

		// 添加到集合映射
		cm.collections[meta.Name] = db
		fmt.Printf("成功加载集合: %s (来自元数据)\n", meta.Name)
	}

	return nil
}

// 保存元数据到文件
func (cm *CollectionManager) saveMetadata() error {
	cm.mu.Lock()
	defer cm.mu.Unlock()

	// 更新元数据中的文档数量
	for i, meta := range cm.metadata.Collections {
		if db, exists := cm.collections[meta.Name]; exists {
			if docCount, err := cm.getDocumentCount(db); err == nil {
				cm.metadata.Collections[i].DocumentCount = docCount
			}
			cm.metadata.Collections[i].LastModified = time.Now()
		}
	}

	// 序列化为JSON
	data, err := json.MarshalIndent(cm.metadata, "", "  ")
	if err != nil {
		return fmt.Errorf("序列化元数据失败: %w", err)
	}

	// 写入文件
	if err := os.WriteFile(cm.metaFile, data, 0644); err != nil {
		return fmt.Errorf("保存元数据文件失败: %w", err)
	}

	return nil
}

// 扫描目录查找集合文件
func (cm *CollectionManager) scanDirectoryForCollections() {
	// 读取目录内容
	files, err := os.ReadDir(cm.baseDir)
	if err != nil {
		fmt.Printf("读取集合目录失败: %v\n", err)
		return
	}

	// 重置元数据
	cm.metadata = model.CollectionsMetadata{
		Collections: []model.CollectionMetadata{},
	}

	// 遍历目录中的所有.sqlite文件
	for _, file := range files {
		if file.IsDir() || !strings.HasSuffix(file.Name(), ".sqlite") {
			continue
		}

		// 跳过临时文件和WAL/SHM文件
		if strings.HasSuffix(file.Name(), ".tmp") ||
			strings.HasSuffix(file.Name(), ".sqlite-wal") ||
			strings.HasSuffix(file.Name(), ".sqlite-shm") {
			continue
		}

		// 提取集合名称（去掉.sqlite后缀）
		collName := strings.TrimSuffix(file.Name(), ".sqlite")

		// 构建完整文件路径
		filePath := filepath.Join(cm.baseDir, file.Name())

		// 获取文件信息（创建时间等）
		fileInfo, err := os.Stat(filePath)
		if err != nil {
			fmt.Printf("获取文件信息失败: %v\n", err)
			continue
		}

		// 创建数据库连接
		db, err := NewSQLiteVectorDB(filePath)
		if err != nil {
			fmt.Printf("加载集合 %s 失败: %v\n", collName, err)
			continue
		}

		// 添加到集合映射
		cm.collections[collName] = db
		fmt.Printf("成功加载集合: %s (通过扫描)\n", collName)

		// 获取文档数量
		docCount, _ := cm.getDocumentCount(db)

		// 添加到元数据
		cm.metadata.Collections = append(cm.metadata.Collections, model.CollectionMetadata{
			Name:          collName,
			FilePath:      filePath,
			CreatedAt:     fileInfo.ModTime(),
			LastModified:  fileInfo.ModTime(),
			DocumentCount: docCount,
		})
	}

	// 保存扫描得到的元数据
	if err := cm.saveMetadata(); err != nil {
		fmt.Printf("保存扫描得到的元数据失败: %v\n", err)
	}
}

// 获取数据库中的文档数量
func (cm *CollectionManager) getDocumentCount(db *SQLiteVectorDB) (int64, error) {
	var count int64
	err := db.db.Get(&count, "SELECT COUNT(*) FROM documents")
	if err != nil {
		return 0, fmt.Errorf("获取文档数量失败: %w", err)
	}
	return count, nil
}

// CreateCollection 创建新的文档集合
func (cm *CollectionManager) CreateCollection(config model.CollectionConfig) (*SQLiteVectorDB, error) {
	cm.mu.Lock()
	defer cm.mu.Unlock()

	// 检查集合名是否已存在
	if _, exists := cm.collections[config.Name]; exists {
		return nil, fmt.Errorf("集合已存在: %s", config.Name)
	}

	// 确定数据库文件路径
	filePath := config.FilePath
	if filePath == "" {
		filePath = filepath.Join(cm.baseDir, fmt.Sprintf("%s.sqlite", config.Name))
	}

	// 创建数据库
	db, err := NewSQLiteVectorDB(filePath)
	if err != nil {
		return nil, err
	}

	// 初始化数据库
	if err := db.InitDB(); err != nil {
		db.Close()
		return nil, err
	}

	// 保存到集合
	cm.collections[config.Name] = db

	// 添加到元数据
	cm.metadata.Collections = append(cm.metadata.Collections, model.CollectionMetadata{
		Name:         config.Name,
		FilePath:     filePath,
		CreatedAt:    time.Now(),
		LastModified: time.Now(),
	})

	// 保存更新后的元数据
	if err := cm.saveMetadata(); err != nil {
		fmt.Printf("保存集合元数据失败: %v\n", err)
		// 继续执行，不中断创建过程
	}

	return db, nil
}

// GetCollection 获取指定名称的集合
func (cm *CollectionManager) GetCollection(name string) (*SQLiteVectorDB, error) {
	cm.mu.RLock()
	defer cm.mu.RUnlock()

	db, exists := cm.collections[name]
	if !exists {
		return nil, fmt.Errorf("集合不存在: %s", name)
	}

	return db, nil
}

// ListCollections 列出所有可用的集合
func (cm *CollectionManager) ListCollections() []string {
	cm.mu.RLock()
	defer cm.mu.RUnlock()

	var collections []string
	for name := range cm.collections {
		collections = append(collections, name)
	}

	return collections
}

// DeleteCollection 删除集合
func (cm *CollectionManager) DeleteCollection(name string) error {
	cm.mu.Lock()
	defer cm.mu.Unlock()

	db, exists := cm.collections[name]
	if !exists {
		return fmt.Errorf("集合不存在: %s", name)
	}

	// 关闭数据库连接
	if err := db.Close(); err != nil {
		return fmt.Errorf("关闭数据库失败: %w", err)
	}

	// 如果是内存模式且有持久化文件，删除文件
	if db.filePath != "" && db.filePath != ":memory:" {
		if err := os.Remove(db.filePath); err != nil {
			return fmt.Errorf("删除数据库文件失败: %w", err)
		}

		// 删除相关的WAL和SHM文件（如果存在）
		os.Remove(db.filePath + "-wal")
		os.Remove(db.filePath + "-shm")
	}

	// 从集合中移除
	delete(cm.collections, name)

	// 从元数据中移除
	for i, meta := range cm.metadata.Collections {
		if meta.Name == name {
			// 移除此项
			cm.metadata.Collections = append(
				cm.metadata.Collections[:i],
				cm.metadata.Collections[i+1:]...,
			)
			break
		}
	}

	// 保存更新后的元数据
	if err := cm.saveMetadata(); err != nil {
		fmt.Printf("保存集合元数据失败: %v\n", err)
		// 继续执行，不中断删除过程
	}

	return nil
}

// CloseAll 关闭所有集合
func (cm *CollectionManager) CloseAll() error {
	cm.mu.Lock()
	defer cm.mu.Unlock()

	var allErrors []string

	for name, db := range cm.collections {
		// 确保先保存数据再关闭
		if err := db.SaveToDisk(); err != nil {
			allErrors = append(allErrors, fmt.Sprintf("保存集合 %s 失败: %v", name, err))
		}

		if err := db.Close(); err != nil {
			allErrors = append(allErrors, fmt.Sprintf("关闭集合 %s 失败: %v", name, err))
		}
	}

	// 清空集合映射
	cm.collections = make(map[string]*SQLiteVectorDB)

	if len(allErrors) > 0 {
		return fmt.Errorf("关闭集合时发生错误: %s", strings.Join(allErrors, "; "))
	}

	return nil
}

// GetCollectionsInfo 获取所有集合的信息，包括文档数量
func (cm *CollectionManager) GetCollectionsInfo() ([]model.CollectionMetadata, error) {
	cm.mu.RLock()
	defer cm.mu.RUnlock()

	// 先更新元数据中的文档数量
	for i, meta := range cm.metadata.Collections {
		if db, exists := cm.collections[meta.Name]; exists {
			if docCount, err := cm.getDocumentCount(db); err == nil {
				cm.metadata.Collections[i].DocumentCount = docCount
			}
		}
	}

	// 返回集合元数据的副本，包含最新的文档数量
	result := make([]model.CollectionMetadata, len(cm.metadata.Collections))
	copy(result, cm.metadata.Collections)
	return result, nil
}

// 修改后的 NewSQLiteVectorDB 函数
func NewSQLiteVectorDB(filePath string) (*SQLiteVectorDB, error) {
	var dsn string
	memoryMode := false

	if filePath == "" || filePath == ":memory:" {
		dsn = ":memory:"
		memoryMode = true
	} else {
		// 确保目录存在
		dir := filepath.Dir(filePath)
		if dir != "." {
			if err := os.MkdirAll(dir, 0755); err != nil {
				return nil, fmt.Errorf("创建目录失败: %w", err)
			}
		}

		dsn = filePath
	}

	sqlDB, err := sqlx.Connect("sqlite3", dsn)
	if err != nil {
		return nil, fmt.Errorf("连接数据库失败: %w", err)
	}

	db := &SQLiteVectorDB{
		db:               sqlDB,
		filePath:         filePath,
		memoryMode:       memoryMode,
		mu:               sync.RWMutex{},
		lastPersistTime:  time.Now(),
		autoSaveInterval: 5 * time.Minute, // 默认5分钟自动保存一次
		modified:         false,
	}

	// 确保在内存模式且有文件路径的情况下启动自动保存
	if memoryMode && filePath != "" && filePath != ":memory:" {
		// 首先尝试从磁盘加载
		if err := db.LoadFromDisk(); err != nil && !os.IsNotExist(err) {
			fmt.Printf("从磁盘加载数据失败: %v", err)
		}

		// 启动自动保存协程
		go db.autoSaveRoutine()
	}

	return db, nil
}

// autoSaveRoutine 自动定期将内存数据保存到磁盘
func (s *SQLiteVectorDB) autoSaveRoutine() {
	ticker := time.NewTicker(s.autoSaveInterval)
	defer ticker.Stop()

	for {
		select {
		case <-ticker.C:
			s.mu.RLock()
			modified := s.modified
			s.mu.RUnlock()

			if modified {
				if err := s.SaveToDisk(); err != nil {
					fmt.Printf("自动保存到磁盘失败: %v\n", err)
				}
			}
		case <-s.stopAutoSave:
			return
		}
	}
}

// LoadFromDisk 从磁盘加载数据到内存数据库
func (s *SQLiteVectorDB) LoadFromDisk() error {
	if !s.memoryMode || s.filePath == ":memory:" {
		return fmt.Errorf("只能在内存模式下从磁盘加载数据")
	}

	// 如果文件不存在，没有需要加载的内容
	if _, err := os.Stat(s.filePath); os.IsNotExist(err) {
		return nil
	}

	// 从磁盘打开数据库连接
	diskDB, err := sqlx.Connect("sqlite3", s.filePath)
	if err != nil {
		return fmt.Errorf("连接磁盘数据库失败: %w", err)
	}
	defer diskDB.Close()

	// 备份数据库内容到内存
	_, err = diskDB.Exec("VACUUM INTO ?", ":memory:")
	if err != nil {
		return fmt.Errorf("从磁盘加载数据到内存失败: %w", err)
	}

	return nil
}

// SaveToDisk 将内存数据库保存到磁盘
func (s *SQLiteVectorDB) SaveToDisk() error {
	s.mu.Lock()
	defer s.mu.Unlock()

	if !s.memoryMode || s.filePath == ":memory:" {
		return nil // 非内存模式或没有指定文件路径，无需保存
	}

	// 创建备份文件
	tmpFile := s.filePath + ".tmp"

	// 将内存数据库保存到临时文件
	_, err := s.db.Exec("VACUUM INTO ?", tmpFile)
	if err != nil {
		return fmt.Errorf("保存到临时文件失败: %w", err)
	}

	// 替换原文件
	if err := os.Rename(tmpFile, s.filePath); err != nil {
		return fmt.Errorf("替换数据库文件失败: %w", err)
	}

	s.lastPersistTime = time.Now()
	s.modified = false

	return nil
}

// InitDB 初始化数据库，创建必要的表
func (s *SQLiteVectorDB) InitDB() error {
	// 创建文档表，存储文本内容和向量
	schema := `
	CREATE TABLE IF NOT EXISTS documents (
		id INTEGER PRIMARY KEY AUTOINCREMENT,
		content TEXT NOT NULL,
		vector BLOB NOT NULL
	);
	`

	_, err := s.db.Exec(schema)
	if err != nil {
		return fmt.Errorf("创建表失败: %w", err)
	}

	// 启用WAL模式提高性能
	_, err = s.db.Exec("PRAGMA journal_mode=WAL;")
	if err != nil {
		return fmt.Errorf("设置WAL模式失败: %w", err)
	}

	// 优化SQLite配置
	_, err = s.db.Exec("PRAGMA synchronous=NORMAL;")
	if err != nil {
		return fmt.Errorf("设置synchronous模式失败: %w", err)
	}

	return nil
}

// AddDocument 添加新文档到数据库
func (s *SQLiteVectorDB) AddDocument(content string, vector []float32) (int64, error) {
	s.mu.Lock()
	defer s.mu.Unlock()

	// 将向量序列化为二进制数据
	vectorBytes, err := json.Marshal(vector)
	if err != nil {
		return 0, fmt.Errorf("序列化向量失败: %w", err)
	}

	query := "INSERT INTO documents (content, vector) VALUES (?, ?)"
	result, err := s.db.Exec(query, content, vectorBytes)
	if err != nil {
		return 0, fmt.Errorf("插入文档失败: %w", err)
	}

	id, err := result.LastInsertId()
	if err != nil {
		return 0, fmt.Errorf("获取插入ID失败: %w", err)
	}

	// 标记为已修改，需要持久化
	s.modified = true

	return id, nil
}

// UpdateDocument 更新文档
func (s *SQLiteVectorDB) UpdateDocument(id int64, content string, vector []float32) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	vectorBytes, err := json.Marshal(vector)
	if err != nil {
		return fmt.Errorf("序列化向量失败: %w", err)
	}

	query := "UPDATE documents SET content = ?, vector = ? WHERE id = ?"
	result, err := s.db.Exec(query, content, vectorBytes, id)
	if err != nil {
		return fmt.Errorf("更新文档失败: %w", err)
	}

	rowsAffected, err := result.RowsAffected()
	if err != nil {
		return fmt.Errorf("获取更新结果失败: %w", err)
	}

	if rowsAffected == 0 {
		return fmt.Errorf("文档不存在: ID=%d", id)
	}

	// 标记为已修改，需要持久化
	s.modified = true

	return nil
}

// DeleteDocument 删除文档
func (s *SQLiteVectorDB) DeleteDocument(id int64) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	query := "DELETE FROM documents WHERE id = ?"
	result, err := s.db.Exec(query, id)
	if err != nil {
		return fmt.Errorf("删除文档失败: %w", err)
	}

	rowsAffected, err := result.RowsAffected()
	if err != nil {
		return fmt.Errorf("获取删除结果失败: %w", err)
	}

	if rowsAffected == 0 {
		return fmt.Errorf("文档不存在: ID=%d", id)
	}

	// 标记为已修改，需要持久化
	s.modified = true

	return nil
}

// DeleteAllDocuments 删除集合中的所有文档
func (s *SQLiteVectorDB) DeleteAllDocuments() (int64, error) {
	s.mu.Lock()
	defer s.mu.Unlock()

	query := "DELETE FROM documents"
	result, err := s.db.Exec(query)
	if err != nil {
		return 0, fmt.Errorf("删除所有文档失败: %w", err)
	}

	rowsAffected, err := result.RowsAffected()
	if err != nil {
		return 0, fmt.Errorf("获取删除结果失败: %w", err)
	}

	// 标记为已修改，需要持久化
	s.modified = true

	return rowsAffected, nil
}

// GetDocument 获取指定ID的文档
func (s *SQLiteVectorDB) GetDocument(id int64) (*model.Document, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	var doc struct {
		ID      int64  `db:"id"`
		Content string `db:"content"`
		Vector  []byte `db:"vector"`
	}

	query := "SELECT id, content, vector FROM documents WHERE id = ?"
	err := s.db.Get(&doc, query, id)
	if err != nil {
		if err == sql.ErrNoRows {
			return nil, fmt.Errorf("文档不存在: ID=%d", id)
		}
		return nil, fmt.Errorf("查询文档失败: %w", err)
	}

	var vector []float32
	if err := json.Unmarshal(doc.Vector, &vector); err != nil {
		return nil, fmt.Errorf("解析向量失败: %w", err)
	}

	return &model.Document{
		ID:      doc.ID,
		Content: doc.Content,
		Vector:  vector,
	}, nil
}

// SearchSimilar 搜索相似文档
func (s *SQLiteVectorDB) SearchSimilar(queryVector []float32, limit int, minScore float32) ([]model.SearchResult, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	// 获取所有文档
	var docs []struct {
		ID      int64  `db:"id"`
		Content string `db:"content"`
		Vector  []byte `db:"vector"`
	}

	query := "SELECT id, content, vector FROM documents"
	err := s.db.Select(&docs, query)
	if err != nil {
		return nil, fmt.Errorf("查询文档失败: %w", err)
	}

	// 计算余弦相似度并排序
	var results []model.SearchResult

	for _, doc := range docs {
		var vector []float32
		if err := json.Unmarshal(doc.Vector, &vector); err != nil {
			return nil, fmt.Errorf("解析向量失败: %w", err)
		}

		// 计算余弦相似度
		score := cosineSimilarity(queryVector, vector)

		// 过滤低于最小分数的结果
		if score >= minScore {
			results = append(results, model.SearchResult{
				Document: model.Document{
					ID:      doc.ID,
					Content: doc.Content,
					Vector:  vector,
				},
				Score: score,
			})
		}
	}

	// 按相似度降序排序
	sort.Slice(results, func(i, j int) bool {
		return results[i].Score > results[j].Score
	})

	// 限制结果数量
	if limit > 0 && len(results) > limit {
		results = results[:limit]
	}

	return results, nil
}

// cosineSimilarity 计算两个向量的余弦相似度
func cosineSimilarity(a, b []float32) float32 {
	if len(a) != len(b) {
		return 0
	}

	var dotProduct float64
	var normA float64
	var normB float64

	for i := 0; i < len(a); i++ {
		dotProduct += float64(a[i] * b[i])
		normA += float64(a[i] * a[i])
		normB += float64(b[i] * b[i])
	}

	if normA == 0 || normB == 0 {
		return 0
	}

	return float32(dotProduct / (math.Sqrt(normA) * math.Sqrt(normB)))
}

// CosineSimilarityExport 导出余弦相似度函数供测试使用
func CosineSimilarityExport(a, b []float32) float32 {
	return cosineSimilarity(a, b)
}

// Close 关闭数据库连接
func (s *SQLiteVectorDB) Close() error {
	s.mu.Lock()
	defer s.mu.Unlock()

	// 如果是内存模式且有文件路径，保存到磁盘
	if s.memoryMode && s.filePath != ":memory:" && s.modified {
		// 停止自动保存协程
		close(s.stopAutoSave)

		// 保存到磁盘
		if err := s.saveToDiskLocked(); err != nil {
			return fmt.Errorf("保存到磁盘失败: %w", err)
		}
	}

	return s.db.Close()
}

// saveToDiskLocked 是SaveToDisk的内部版本，假设已经持有锁
func (s *SQLiteVectorDB) saveToDiskLocked() error {
	if !s.memoryMode || s.filePath == ":memory:" {
		return nil // 非内存模式或没有指定文件路径，无需保存
	}

	// 创建备份文件
	tmpFile := s.filePath + ".tmp"

	// 将内存数据库保存到临时文件
	_, err := s.db.Exec("VACUUM INTO ?", tmpFile)
	if err != nil {
		return fmt.Errorf("保存到临时文件失败: %w", err)
	}

	// 替换原文件
	if err := os.Rename(tmpFile, s.filePath); err != nil {
		return fmt.Errorf("替换数据库文件失败: %w", err)
	}

	s.lastPersistTime = time.Now()
	s.modified = false

	return nil
}

// ListDocuments 获取文档列表，返回指定范围的文档和总文档数
func (db *SQLiteVectorDB) ListDocuments(limit, offset int) ([]model.Document, int64, error) {
	// 获取文档总数
	var total int64
	err := db.db.Get(&total, "SELECT COUNT(*) FROM documents")
	if err != nil {
		return nil, 0, fmt.Errorf("获取文档数量失败: %w", err)
	}

	// 获取分页文档
	var docs []model.Document
	err = db.db.Select(&docs, `
        SELECT id, content FROM documents 
        ORDER BY id 
        LIMIT ? OFFSET ?
    `, limit, offset)

	if err != nil {
		return nil, 0, fmt.Errorf("获取文档列表失败: %w", err)
	}

	return docs, total, nil
}
